#!/usr/bin/python
"""

   parse_intron_exon_from_fa.py

"""

from optparse import OptionParser
import sys, os
import re
import StringIO


import sys

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(os.path.abspath(os.path.join(exe_path, "..", "python_modules")))

from general_util import die_error
from translate_by_codon import translate

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(os.path.abspath(os.path.join(exe_path,"..", "python_modules")))

import general_util
from custom_log import custom_log
from general_util import die_error
from exon_structure import gene_exons




parser = OptionParser(version="%prog 1.0")
parser.add_option("-d", "--dna_file", dest="dna_file",
                  metavar="FILE",                       
                  type="string",
                  help="(Mandatory) name and path of dna file")


#
#   Output files
# 
parser.add_option("-C", "--compare_cdna_by_exons", dest="compare_cdna_by_exons_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold aligned cDNA broken up into "+
                        "exons.")


parser.add_option("-g", "--genomic_by_exons", dest = "genomic_by_exons_file",
                  metavar="FILE",                       
                  type="string",
                  help="name and path of file for genomic DNA sequence broken up into "\
                        "introns and exons.")
parser.add_option("-e", "--exons", dest="exons_file",
                  metavar="FILE",
                  type="string",
                  help="name and path of file for exons.")
parser.add_option("-c", "--cds_in_exons", dest="cds_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold cDNA broken up into exons.")

parser.add_option("-p", "--prot_in_exons", dest="prot_in_exons_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold translated peptide sequence "\
                        "broken up into exons.")
parser.add_option("-P", "--peptide", dest="peptide_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold continguous translated peptide sequence.")

parser.add_option("-l", "--lengths_and_splice", dest = "lengths_and_splice_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to hold exon and intron lengths and splice_sites.")
parser.add_option("-s", "--summary", dest="summary_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of file to total lengths of different GFFs.")
parser.add_option("-w", "--wrap_genomic", dest = "wrap_genomic",
                  action="store_true", default=False,
                  help="line wrap (main) genomic output file every 60 residues")
parser.add_option("-N", "--gene_name", dest = "gene_name",
                  metavar="NAME", 
                  type="string",
                  default = '',
                  help="NAME of gene on fasta description line")



file_parameters = """
dna_file
genomic_by_exons_file
exons_file
cds_file
prot_in_exons_file
peptide_file
lengths_and_splice_file
summary_file""".split()

#
#   everything defaults to None
# 
for parameter in file_parameters :
    parser.set_defaults(**{parameter: None}) 




# get help string
f =StringIO.StringIO()
parser.print_help(f)
helpstr = f.getvalue()
(options, args) = parser.parse_args()


# 
# mandatory options
# 
for parameter in ["dna_file"] :
    if getattr(options, parameter) == None:
        die_error("--dna_file is mandatory. "\
                "Please specify this files.\n\n" + helpstr)


#
#   open output files
# 
def open_or_stdout(file_name):
    if file_name == None: 
        return None
    if file_name == '-':
        return sys.stdout
    else:
        return open(file_name,  "w")

exons_file                  = open_or_stdout(options.exons_file         )
prot_in_exons_file          = open_or_stdout(options.prot_in_exons_file )
peptide_file                = open_or_stdout(options.peptide_file )
lengths_and_splice_file     = open_or_stdout(options.lengths_and_splice_file)
summary_file                = open_or_stdout(options.summary_file)
genomic_by_exons_file       = open_or_stdout(options.genomic_by_exons_file)
compare_cdna_by_exons_file  = open_or_stdout(options.compare_cdna_by_exons_file)
cds_file                    = open_or_stdout(options.cds_file)


if len(options.gene_name):
    options.gene_name += ' '

dnalines = open(options.dna_file)
dna, gene = gene_exons.load_from_fasta(dnalines)

from itertools import izip
#
#   helper function to write out exons/introns
# 
def write_intron_exons (fileH, gene, iter):
    fileH.write("#" + options.gene_name + gene.get_summary() + "\n")
    for description, seq in iter:
        fileH.write( ">" + description + "\n")
        # wrap lines
        if options.wrap_genomic:
            fileH.write(line_wrap(seq, 60) + "\n")
        else:
            fileH.write(seq + "\n")

def write_exons (fileH, gene, iter):
    fileH.write("#" + options.gene_name + gene.get_summary() + "\n")
    for description, seq in izip(gene.iter_descriptions(), iter):
        fileH.write( ">" + description + "\n")
        # wrap lines
        if options.wrap_genomic:
            fileH.write(line_wrap(seq, 60) + "\n")
        else:
            fileH.write(seq + "\n")
    

            
# 
# splice sites and lengths
# 
if lengths_and_splice_file:
    line_format = "%-9s\t" * 7  + "\n"
    lengths_and_splice_file.write(line_format % gene.tabulated_summary_columns())
    for cols in gene.tabulated_summary(dna):
        lengths_and_splice_file.write(line_format % tuple(cols))

#
#   summary
# 
if summary_file:
    summary_file.write(gene.get_summary() + "\n")

#
#   genomic intron exon
# 
if genomic_by_exons_file:
    write_intron_exons(genomic_by_exons_file, gene, gene.iter_intron_exons(dna, 500))
        
#
#   all exons
# 
if exons_file:
    write_exons(exons_file, gene, gene.iter_exons(dna))
        


#
#   translations by exons
# 
if prot_in_exons_file:
    write_exons(prot_in_exons_file, gene, gene.iter_translation(dna, True))

if peptide_file:
    peptide_file.write(">" + options.gene_name + gene.get_summary() + "\n")
    for seq in gene.iter_translation(dna, False):
        if not len(seq):
            continue
        # wrap lines
        if options.wrap_genomic:
            peptide_file.write(line_wrap(seq, 60) + "\n")
        else:
            peptide_file.write(seq + "\n")

#
#   CDS by exons
# 
if cds_file:
    write_exons(cds_file, gene, gene.iter_cds(dna))
                    



#if compare_cdna_by_exons_file:
#    gene.print_matching_exon_sequence(dna, query_dna, compare_cdna_by_exons_file)

